
import PyPDF2
import sys

class PDFDirGenerator:
    
    def __init__(self, params=None, isPrint = True) -> None:
        
        self.pdf_input_path = ''
        self.pdf_output_path = ''
        self.dir_path = ''
        self.prefix = '\t'
        self.delimiter = ' '
        self.offset = 0
        
        self.dir_parents = [None]
    
        self.isPrint = isPrint
    
        if params:
            
            if 'PDF_INPUT_PATH' in params:
                self.pdf_input_path = params['PDF_INPUT_PATH']
            if 'PDF_OUTPUT_PATH' in params:
                self.pdf_output_path = params['PDF_OUTPUT_PATH']
            if 'DIR_INPUT_PATH' in params:
                self.dir_path = params['DIR_INPUT_PATH']
            if 'PREFIX' in params:
                self.prefix = params['PREFIX']
                if self.prefix == "\\t":
                    self.prefix = '\t'
            if 'DELIMITER' in params:
                self.delimiter = params['DELIMITER']
            if 'OFFSET' in params:
                self.offset = int(params['OFFSET'])
        
            
    def reset(self, params):
        
        if params:
            
            if 'PDF_INPUT_PATH' in params:
                self.pdf_input_path = params['PDF_INPUT_PATH']
            if 'PDF_OUTPUT_PATH' in params:
                self.pdf_output_path = params['PDF_OUTPUT_PATH']
            if 'DIR_INPUT_PATH' in params:
                self.dir_path = params['DIR_INPUT_PATH']
            if 'PREFIX' in params:
                self.prefix = params['PREFIX']
                if self.prefix == "\\t":
                    self.prefix = '\t'
            if 'DELIMITER' in params:
                self.delimiter = params['DELIMITER']
            if 'OFFSET' in params:
                self.offset = int(params['OFFSET'])
    
    def getLevel(self, line:str):
        level = 0
        for ch in line:
            if ch == self.prefix:
                level += 1
            else:
                break
        # BUG 级别必须从1开始算起
        return level + 1
    
    def getInfo(self):
        
        info_str = \
            " PDF input path: %s\n"\
            " PDF output path: %s\n"\
            " Dir path: %s\n"\
            " Offset: %d\n"\
            " Prefix: %s\n"\
            " Delimiter: %s"\
            % (self.pdf_input_path,
               self.pdf_output_path,
               self.dir_path,
               self.offset,
               self.prefix,
               self.delimiter
            )
        
        return info_str
    
    def getDividePos(self, line):
        # divide_pos = line.rfind(self.delimiter) # BUG 希望 \t 空格都包含在内
        if self.delimiter == ' ':
            divide_pos = max(line.rfind(' '), line.rfind('\t'))
        else:
            divide_pos = line.rfind(self.delimiter)
        return divide_pos
    
    def run(self):
        
        res = False
        
        if self.isPrint:
            print("--------------------------- Adding the bookmark ---------------------------")
            print("%s" % self.getInfo())
            print("---------------------------------------------------------------------------")
        
        
        with open(self.dir_path, 'r', encoding='utf-8') as fin:
            
            pdf_reader = PyPDF2.PdfFileReader(self.pdf_input_path)
            pdf_writer = PyPDF2.PdfFileWriter()
            
            pdf_writer.cloneDocumentFromReader(pdf_reader)
            
            lines = fin.readlines()
            
            for i, oline in enumerate(lines):
                line = oline.strip()
                level = self.getLevel(oline)
                divide_pos = self.getDividePos(line)
                page = int(line[divide_pos+1:]) + self.offset
                # 添加书签
                if level >= len(self.dir_parents):
                    self.dir_parents.append(None)
                self.dir_parents[level] = pdf_writer.addBookmark(
                    line[0:divide_pos], page-1, self.dir_parents[level-1]
                )
                
                if self.isPrint:
                    print(" [%3d/%3d finished] level: %d, title: %s, page: %d"
                          % (
                              i+1, len(lines), level, line[0:divide_pos], page
                          ))
        
            # 设置输出目录                        
            if len(self.pdf_output_path) == 0:
                dot_pos = self.pdf_input_path.rfind('.')
                self.pdf_output_path = self.pdf_input_path[0:dot_pos] + '(书签).pdf'
            else:
                output_dot_pos = self.pdf_output_path.rfind('.')
                output_base_pos = self.pdf_output_path.rfind('\\')
                if output_dot_pos != -1 and output_base_pos < output_dot_pos:
                    # 给定文件名
                    pass 
                else:
                    # 给定文件夹
                    dot_pos = self.pdf_input_path.rfind('.')
                    base_pos = self.pdf_input_path.rfind('\\')
                    base_name = self.pdf_input_path[base_pos+1:dot_pos]
                    if self.pdf_output_path[-1] == '\\':
                        self.pdf_output_path = self.pdf_output_path + base_name + '(书签).pdf'
                    else:
                        self.pdf_output_path = self.pdf_output_path + '\\' + base_name + '(书签).pdf'
            
            with open(self.pdf_output_path, 'wb') as fout:
                pdf_writer.write(fout)
                res = True

                if self.isPrint:
                    print("---------------------------------------------------------------------------")
                    print(" Save: %s" % self.pdf_output_path)
                    print("---------------------------------- Done! ----------------------------------")
 
        return res
                

if __name__ == '__main__':
    
    # 运行指令
    # python .\pdfDirGenerator.py <dir_input_path> <pdf_input_path> [<-option> <params>]
    
    argc = len(sys.argv)
    params = {
        'DIR_INPUT_PATH' : sys.argv[1],
        'PDF_INPUT_PATH' : sys.argv[2]}
    for i in range(3, argc, 2):
        key = sys.argv[i]
        value = sys.argv[i+1]
        if key == '-o':
            key = 'PDF_OUTPUT_PATH'
        elif key == '-p':
            key = 'PREFIX'
        elif key == '-d':
            key = 'DELIMITER'
        elif key == '-O':
            key = 'OFFSET'
        else:
            continue
        params[key] = value
        
    # params = {
    #     "PDF_INPUT_PATH" : 'testpdf.pdf',
    #     "PDF_OUTPUT_PATH" : 'C:\\Users\\q2799\\Project\\PDFDirGenerator2\\test\\',
    #     "DIR_INPUT_PATH" : 'dir_format.txt',
    #     "PREFIX" : '\\t',
    #     "DELIMITER" : ' ',
    #     "OFFSET" : 12
    # }
    
    pdg = PDFDirGenerator(params)
    pdg.run()